import json
import os

import matplotlib.pyplot as plt
from cached_property import cached_property
from docx import Document
from docx.enum.style import WD_STYLE_TYPE
from docx.shared import Pt

from section import Section, LowerLevelSection
from settings import RESULTS_DIR
from stack import (InterchangeStack, Stack, StraightStack, TunnelStack, ClearAreaStack,
                   VerticalStack, recreate_stack_number)
from worksheet import StraightWorkSheet, VerticalWorkSheet


class Worker:
    def __init__(self, s_sheet_filenames, v_sheet_filenames, reverse=False, other_stacks=None, output_name=None,
                 config=None, level=1, **kwargs):
        self.reverse = reverse
        self.output_name = output_name
        self.s_sheets = [StraightWorkSheet(filename, output_name=output_name, reverse=reverse)
                         for filename in s_sheet_filenames]
        self.v_sheets = [VerticalWorkSheet(filename, output_name=output_name, reverse=reverse)
                         for filename in v_sheet_filenames]
        self.raw_other_stacks = other_stacks
        self.config = config
        self.level = level
        os.makedirs(self.path, exist_ok=True)

    @cached_property
    def section_class(self):
        return Section if self.level == 1 else LowerLevelSection

    @cached_property
    def other_stacks(self):
        tunnel_stacks, interchange_stacks, clear_area_stacks = self.parse_raw_stack_list()
        if len(tunnel_stacks+interchange_stacks+clear_area_stacks) == 0:
            return []
        tunnel_stacks = self.preprocess_tunnel_stacks(
            tunnel_stacks)  # 短隧道与间距小于两百米的隧道处理

        all_stacks = interchange_stacks + tunnel_stacks
        all_stacks.sort(key=lambda s: s.meters)

        processed_stacks = []
        in_tunnel = False
        # 先将隧道段中的其他桩全部去掉
        for stack in all_stacks:
            if not in_tunnel:
                in_tunnel = isinstance(stack, TunnelStack)
                processed_stacks.append(stack)
            else:
                in_tunnel = not isinstance(stack, TunnelStack)
                if not in_tunnel:
                    processed_stacks.append(stack)

        results = []

        for i in range(len(processed_stacks)):
            stack = processed_stacks[i]
            results.append(stack)
            # 碰到不是成对出现的互通立交桩就创建相应的新桩
            if isinstance(stack, InterchangeStack):
                is_start = stack.is_start
                adjacent_stack: Stack = processed_stacks[i +
                                                         1] if is_start else processed_stacks[i - 1]
                if not isinstance(adjacent_stack, InterchangeStack):
                    offset = -0.1 if is_start else 0.1
                    results.append(adjacent_stack.new_with_offset(
                        offset, stack_type='interchange'))
        results.extend(clear_area_stacks)
        results.sort(key=lambda s: s.meters)

        return results

    @cached_property
    def stacks(self):
        stacks = self.s_sheets_stacks + \
            self.v_sheets_stacks + self.other_stacks
        stacks = self.preprocess_stacks(stacks)

        stacks.sort(key=lambda stack: stack.meters, reverse=self.reverse)
        return stacks

    @cached_property
    def s_sheets_stacks(self):
        # 上一个文件最后一个路段的终点到下一个文件第一个大于它的起点
        sheets = self.s_sheets
        results = [*sheets[0].stacks]
        for i, sheet in enumerate(sheets[:-1]):
            next_sheet = sheets[i+1]
            for j in range(len(next_sheet.stacks)):
                if next_sheet.stacks[j] >= sheet.stacks[-1]:
                    break
            j = j if j % 2 == 0 else j+1
            results.extend(next_sheet.stacks[j:])
        return results

    @cached_property
    def v_sheets_stacks(self):
        # 上一个文件的最后一个桩到下一个文件第一个大于它的桩
        sheets = self.v_sheets
        results = [*sheets[0].stacks]
        for i, sheet in enumerate(sheets[:-1]):
            next_sheet = sheets[i+1]
            for j in range(len(next_sheet.stacks)):
                if next_sheet.stacks[j] >= sheet.stacks[-1]:
                    break
                results.extend(next_sheet.stacks[j:])
        return results

    def preprocess_tunnel_stacks(self, tunnel_stacks):
        _tunnel_stacks = []

        in_tunnel = False
        i = 0

        # 隧道间距小于两百米，就视为一段隧道
        while i < len(tunnel_stacks):
            if in_tunnel and i + 1 != len(tunnel_stacks):
                if abs(tunnel_stacks[i + 1].meters - tunnel_stacks[i].meters) < 200:
                    i += 2
                else:
                    _tunnel_stacks.append(tunnel_stacks[i])
                    i += 1
                    in_tunnel = False
                    continue
            else:
                _tunnel_stacks.append(tunnel_stacks[i])
                in_tunnel = not in_tunnel
                i += 1

        # 判断是否是短隧道
        for i in range(0, len(_tunnel_stacks), 2):
            if abs(_tunnel_stacks[i + 1].meters - _tunnel_stacks[i].meters) < 200:
                _tunnel_stacks[i + 1].is_short = True
                _tunnel_stacks[i].is_short = True
            else:
                _tunnel_stacks[i + 1].is_short = False
                _tunnel_stacks[i].is_short = False
        return _tunnel_stacks

    def parse_raw_stack_list(self):
        if self.raw_other_stacks is None:
            return [], [], []

        tunnel_stacks = []
        interchange_stacks = []
        clear_area_stacks = []
        for stack in self.raw_other_stacks:
            type_ = stack['类型']
            if type_ == '路侧净区':
                clear_area_stacks.append(ClearAreaStack(stack['桩号'], width=stack['宽度']))
                continue
            start = stack['起点']
            end = stack.get('止点') or stack['终点']
            name = stack.get('段名')
            if type_ == '隧道':
                tunnel_stacks.append(TunnelStack(
                    start, name=name, is_start=True))
                tunnel_stacks.append(TunnelStack(
                    end, name=name, is_start=False))
            elif type_ == '互通立交':
                interchange_stacks.append(
                    InterchangeStack(start, name=name, is_start=True))
                interchange_stacks.append(
                    InterchangeStack(end, name=name, is_start=False))

        tunnel_stacks.sort(key=lambda stack: stack.meters)
        interchange_stacks.sort(key=lambda stack: stack.meters)
        clear_area_stacks.sort(key=lambda stack: stack.meters)

        return tunnel_stacks, interchange_stacks, clear_area_stacks

    def preprocess_stacks(self, stacks):
        """
        去除隧道和互通立交桩之间的桩，并且在隧道入口桩前面200米和隧道出口桩后面100米创建新桩
        """
        stacks.sort(key=lambda stack: stack.meters)
        other_stack_flag = False
        _stacks = []
        for stack in stacks:
            if not other_stack_flag:
                if isinstance(stack, TunnelStack):  # 隧道入口桩
                    # 创建一个隧道入口前200米的新桩 短隧道不创建
                    other_stack_flag = True
                    if not stack.is_short:
                        new_stack = Stack(recreate_stack_number(
                            stack.stack_number, stack.meters - 200))
                        _stacks.append(new_stack)
                    _stacks.append(stack)
                elif isinstance(stack, InterchangeStack):
                    other_stack_flag = True
                    _stacks.append(stack)
                else:
                    _stacks.append(stack)
            else:
                if isinstance(stack, TunnelStack):  # 隧道出口桩
                    # 创建一个隧道出口后100米的新桩 短隧道不创建 有互通立交的也不创建
                    other_stack_flag = False
                    if not stack.is_short:
                        new_stack = Stack(recreate_stack_number(
                            stack.stack_number, stack.meters + 100))
                        _stacks.append(stack)
                        _stacks.append(new_stack)
                    else:
                        _stacks.append(stack)
                elif isinstance(stack, InterchangeStack):
                    other_stack_flag = False
                    _stacks.append(stack)
        _stacks.sort(key=lambda s: s.meters)

        return _stacks

    @cached_property
    def sections(self):
        stacks = self.stacks
        slope_degree = None
        radius = None
        sections = []
        in_tunnel = False
        in_interchange = False
        influence_coefficient = 1
        for i, stack in enumerate(stacks):
            if i == len(stacks) - 1:
                break
            if isinstance(stack, StraightStack):
                radius = stack.radius
            elif isinstance(stack, VerticalStack):
                slope_degree = stack.slope_degree
            elif isinstance(stack, ClearAreaStack):
                influence_coefficient = stack.influence_coefficient
            if slope_degree is not None and radius is not None:
                if isinstance(stack, TunnelStack) and not in_tunnel:
                    section = self.section_class(radius, slope_degree, stack,
                                                 stacks[i + 1], _type='tunnel')
                    in_tunnel = True
                elif isinstance(stack, InterchangeStack) and not in_interchange:
                    section = self.section_class(radius, slope_degree, stack,
                                                 stacks[i + 1], _type='interchange')
                    in_interchange = True
                else:
                    section = self.section_class(
                        radius, slope_degree, stack, stacks[i + 1], influence_coefficient=influence_coefficient)
                sections.append(section)
        for i, section in enumerate(sections):
            if i == 0:
                section.previous_section = None
            else:
                section.previous_section = sections[i - 1]
            try:
                section.latter_section = sections[i + 1]
            except IndexError:
                section.latter_section = None
        for section in sections:
            section.compute()
        return sections

    def plot(self):
        sections = self.sections

        plt.figure()
        x = [section.start.meters for section in sections]
        y = [section.v_in_l for section in sections]
        y2 = [section.v_diff_l for section in sections]
        plt.plot(x, y, 'co--', label='小客车速度')
        plt.plot(x, y2, 'ro-.', label='小客车速度差')
        plt.title(self.project_name)
        plt.legend(loc="best")
        plt.savefig(os.path.join(self.path, '小客车.png'),
                    dpi=500, bbox_inches='tight')

        plt.figure()
        x = [section.start.meters for section in sections]
        y = [section.v_in_b for section in sections]
        y2 = [section.v_diff_b for section in sections]
        plt.plot(x, y, 'co--', label='大货车速度')
        plt.plot(x, y2, 'ro-.', label='大货车速度差')
        plt.title(self.project_name)
        plt.legend(loc="best")
        plt.savefig(os.path.join(self.path, '大货车.png'),
                    dpi=500, bbox_inches='tight')

    @cached_property
    def project_name(self):
        return self.s_sheets[0].project_name

    @cached_property
    def path(self):
        dirname = '下行' if self.reverse else '上行'
        output_name = self.project_name.split(
        )[0] if self.output_name is None else self.output_name
        return os.path.join(RESULTS_DIR, output_name, dirname)

    def save_results(self):
        sections = [section.serialize() for section in self.sections]
        document = Document()
        style = document.styles.add_style('table', WD_STYLE_TYPE.TABLE)
        font = style.font
        font.name = '宋体'
        font.size = Pt(4)
        headers = [
            '单元起点',
            '单元止点',
            '单元类型',
            '半径',
            '坡度',
            '长度',
            '起点小客车速度',
            '止点小客车速度',
            '起点大货车速度',
            '止点大货车速度',
            '小客车速度差',
            '大货车速度差',
            '小客车速度梯度',
            '大货车速度梯度'
        ]
        table = document.add_table(rows=1, cols=len(headers), style='table')
        table.autofit = True,
        header_cells = table.rows[0].cells
        for i, header in enumerate(headers):
            header_cells[i].text = header

        for section in sections:
            row_cells = table.add_row().cells
            row_cells[0].text = section['start']['stack_number']
            row_cells[1].text = section['end']['stack_number']
            row_cells[2].text = section['type']
            row_cells[3].text = str(section['radius'])
            row_cells[4].text = '{:.2f}'.format(section['slope_degree'])
            row_cells[5].text = str(section['length'])
            row_cells[6].text = '{:.2f}'.format(section['v_in_l'])
            row_cells[7].text = '{:.2f}'.format(section['v_out_l'])
            row_cells[8].text = '{:.2f}'.format(section['v_in_b'])
            row_cells[9].text = '{:.2f}'.format(section['v_out_b'])
            row_cells[10].text = '{:.2f}'.format(section['v_diff_l'])
            row_cells[11].text = '{:.2f}'.format(section['v_diff_b'])
            row_cells[12].text = '{:.2f}'.format(section['v_gradient_l'])
            row_cells[13].text = '{:.2f}'.format(section['v_gradient_b'])

        document.add_picture(os.path.join(self.path, '小客车.png'))
        document.add_picture(os.path.join(self.path, '大货车.png'))

        document.save(os.path.join(self.path, 'results.docx'))

        with open(os.path.join(self.path, 'results.json'), mode='w') as file:
            json.dump(sections, file)

    def run(self):
        print(f'{self.project_name[:10]}···:正解析')
        self.plot()
        self.save_results()
        print(f'{self.project_name[:10]}···:解析完毕')
